#!/usr/bin/env python3
"""
Author: Haoran Peng
Email: gavinsweden@gmail.com
"""

import sys
import time
from copy import deepcopy

import cv2
import numpy as np
import yaml

sys.path.append(".")

from cbs_mapf.assigner import *
from cbs_mapf.planner import Planner

assign_flag = 0
pathnum = 1


class Simulator:
    def __init__(self, load_path=None):
        # Set up a white 1080p canvas
        self.canvas = np.ones((5000, 5000, 3), np.uint8) * 255
        # Draw the rectangluar obstacles on canvas
        self.draw_rect(np.array([np.array(v) for v in RECT_OBSTACLES.values()]))

        # Transform the vertices to be border-filled rectangles
        static_obstacles = self.vertices_to_obsts(RECT_OBSTACLES)
        if load_path is None:
            # Call cbs-mapf to plan
            self.planner = Planner(GRID_SIZE, ROBOT_RADIUS, static_obstacles, NUM_AG)
            before = time.time()
            self.path = self.planner.plan(
                START,
                GOAL,
                ASSIGN,
                assign=sequential_assign,
                max_iter=2000,
                low_level_max_iter=1000,
                debug=True,
            )
            after = time.time()
            print("Time elapsed:", "{:.4f}".format(after - before), "second(s)")

            # Assign each agent a colour
            self.colours = self.assign_colour(len(self.path))

            # Put the path into dictionaries for easier access
            d = dict()
            for i, path in enumerate(self.path):
                self.draw_path(self.canvas, path, i)  # Draw the path on canvas
                d[i] = path
            self.path = d
        else:
            self.path = path
            self.colours = self.assign_colour(len(self.path))

    """
    Transform opposite vertices of rectangular obstacles into obstacles
    """

    @staticmethod
    def vertices_to_obsts(obsts):
        def drawRect(v0, v1):
            o = []
            base = abs(v0[0] - v1[0])
            side = abs(v0[1] - v1[1])
            for xx in range(0, base, 30):
                o.append((v0[0] + xx, v0[1]))
                o.append((v0[0] + xx, v0[1] + side - 1))
            o.append((v0[0] + base, v0[1]))
            o.append((v0[0] + base, v0[1] + side - 1))
            for yy in range(0, side, 30):
                o.append((v0[0], v0[1] + yy))
                o.append((v0[0] + base - 1, v0[1] + yy))
            o.append((v0[0], v0[1] + side))
            o.append((v0[0] + base - 1, v0[1] + side))
            return o

        static_obstacles = []
        for vs in obsts.values():
            static_obstacles.extend(drawRect(vs[0], vs[1]))
        return static_obstacles

    """
    Randomly generate colours
    """

    @staticmethod
    def assign_colour(num):
        def colour(x):
            x = hash(str(x + 42))
            return (x & 0xFF, (x >> 8) & 0xFF, (x >> 16) & 0xFF)

        colours = dict()
        for i in range(num):
            colours[i] = colour(i)
        return colours

    def draw_rect(self, pts_arr: np.ndarray) -> None:
        for pts in pts_arr:
            cv2.rectangle(
                self.canvas, tuple(pts[0]), tuple(pts[1]), (0, 0, 255), thickness=3
            )

    def draw_path(self, frame, xys, i):
        for x, y in xys:
            cv2.circle(frame, (int(x), int(y)), 10, self.colours[i], -1)

    """
    Press any key to start.
    Press 'q' to exit.
    """

    def start(self):
        cv2.namedWindow("frame", cv2.WINDOW_NORMAL)

        # 获取屏幕分辨率
        screen_width = 1920  # 替换为你的屏幕宽度
        screen_height = 1080  # 替换为你的屏幕高度

        # 设置窗口大小为屏幕的一部分
        window_width = int(screen_width * 0.8)
        window_height = int(screen_height * 0.8)

        cv2.resizeWindow("frame", window_width, window_height)

        wait = True
        try:
            i = 0
            while True:
                frame = deepcopy(self.canvas)
                for id_ in self.path:
                    x, y = tuple(self.path[id_][i])
                    cv2.circle(frame, (x, y), ROBOT_RADIUS - 5, self.colours[id_], 5)

                cv2.imshow("frame", frame)
                # if wait:
                #     cv2.waitKey(0)
                #     wait = False
                k = cv2.waitKey(100) & 0xFF
                if k == ord("q"):
                    break
                if k == ord("r"):
                    i = 0  # 重来一遍
                i += 1
        except Exception:
            cv2.waitKey(0)
            cv2.destroyAllWindows()


def load_scenario(fd, source_dir):
    with open(fd, "r") as f:
        global GRID_SIZE, ROBOT_RADIUS, RECT_OBSTACLES, START, GOAL
        data = yaml.load(f, Loader=yaml.FullLoader)
        GRID_SIZE = data["GRID_SIZE"]
        ROBOT_RADIUS = data["ROBOT_RADIUS"]
        RECT_OBSTACLES = data["RECT_OBSTACLES"]
        # START = data["START"]
        # GOAL = data["GOAL"]
        load_config(source_dir)


import json


def load_config(fd):
    global START, GOAL, ASSIGN, NUM_AG
    with open(fd + "/TaskInfo.json", "r", encoding="utf-8") as file:
        TaskInfo = json.load(file)
    with open(fd + "/USVPosInit.json", "r", encoding="utf-8") as file:
        USVPosInit = json.load(file)
    usv_pos = [(item[1] + 200, item[2] + 400) for item in USVPosInit["无人艇初始位置"]]

    formation = TaskInfo["agents"][1]["formation"]
    center = TaskInfo["agents"][1]["center"]
    goal_pos = []
    NUM_AG = len(formation)
    for item in formation:
        goal_pos.append(
            (
                center[0]
                - 2000
                + round(item["distance"] * np.cos(np.deg2rad(90 - item["angle"]))),
                center[1]
                + 400
                + round(item["distance"] * np.sin(np.deg2rad(90 - item["angle"]))),
            )
        )
    GOAL = goal_pos
    START = usv_pos
    ASSIGN = formation[0]["usvId"]  # 0则随机分配
    print("START:", START)
    print("GOAL:", GOAL)
    print("ASSIGN:", ASSIGN)


def show_pos(pos):
    cv2.namedWindow("frame", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("frame", (500, 500))
    frame = np.ones((5000, 5000, 3), np.uint8) * 255
    for x, y in pos:
        cv2.circle(frame, (x, y), ROBOT_RADIUS - 5, (0, 0, 0), 5)
    cv2.imshow("frame", frame)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def process_scenario(num):
    source_dir = "test/test" + str(num)
    source_dir = "test/conflict"
    load_scenario("visualization/scenario2.yaml", source_dir)
    # show_pos(START)
    # show_pos(GOAL)
    r = Simulator()
    r.start()
    if len(r.path) != 0:
        np.save("path_dir/conflict_path.npy", r.path)
        np.save("path_dir/conflict_agents_list.npy", r.planner.agents_list)
        # np.save("path_dir/path" + str(num) + ".npy", r.path)
        # np.save("path_dir/agents_list" + str(num) + ".npy", r.planner.agents_list)


if __name__ == "__main__":
    # From command line, call:
    # python3 visualizer.py scenario1.yaml
    # load_scenario(sys.argv[1])
    # 使用 argparse 解析命令行参数
    # parser = argparse.ArgumentParser(description="Process some scenarios.")
    # parser.add_argument(
    #     "--num", type=int, required=True, help="The number of scenarios to process"
    # )
    # args = parser.parse_args()
    num = 1
    # 使用 ThreadPoolExecutor 并行处理任务
    process_scenario(num)
    # with concurrent.futures.ThreadPoolExecutor() as executor:
    #     futures = [
    #         executor.submit(process_scenario, num) for num in range(num, num + 1)
    #     ]
    #     concurrent.futures.wait(futures)
